#! /usr/bin/env ruby

MustBeRoot = true

class Program
  def initialize
    run if MustBeRoot and root?
  rescue Interrupt
  rescue Exception => e
    puts e
    puts e.backtrace
  ensure
    cleanup
  end

  def options
    @options ||= parse_args
  end
end

class Main < Program
  def initialize
    options[:host]
    super
  end

  def run
    cacerts = `find #{ENV['JAVA_HOME'] || '.'} -name cacerts`
    options[:java_keystore] = if File.exists? cacerts
      cacerts
    else
      '/System/Library/Java/Support/CoreDeploy.bundle/Contents/Home/lib/security/cacerts'
    end
    options[:local_keystore] ||= 'LATEST.jks'

    if options[:host]
      host = options[:host]
      port = options[:port] || 443
      a = fetch_cert host, port
      cert_file = host + '.crt'
      File.open(cert_file, 'w') { |f| f << a }
      puts create_keystore(host)
      puts import_keystore(host)
      puts delete_alias(host) if alias_exists?(host)
      puts add_alias(host)
      puts confirm_alias(host)
    else
      for f in Dir['*.crt']
        host = f[/(.*)\..*/, 1]
        jks = host + '.jks'
        puts create_keystore(host) unless File.exists?(jks)
        puts import_keystore(host)
        puts delete_alias(host) if alias_exists?(host)
        puts add_alias(host)
        puts confirm_alias(host)
      end
    end
  end

  def cleanup
    if options[:host]
      host = options[:host]
      cert_file = host + '.crt'
      jks = host + '.jks'
      delete_keystore(host) if File.exists?(jks)
      File.delete cert_file
    end
  end
end

def parse_args
  require 'optparse'

  script_name = File.basename __FILE__

  options = {}
  OptionParser.new do |opt|
    opt.banner = "Usage: #{script_name} [host[:port]] [options...]\n\n"
    opt.on('-h', '--host=host', 'Hostname of the server from which the SSL certificate will be imported.') { |v| options[:host] = v }
    opt.on('-p', '--port=port', 'Port of the server from which the SSL certificate will be imported.') { |v| options[:port] = v }
    opt.on('-f', '--file=file', 'The file containing the SSL certificate to be imported.') { |v| options[:cert] = v }
    opt.on('-k', '--keystore=keystore', 'The local keystore file into which an alias of the certificate will be inserted.') { |v| options[:local_keystore] = v }
    opt.on('-d', '--debug', 'Do you want to see shell commands printed out during execution?') { |v| options[:debug] = true }
    opt.on_tail("-?", "--help", "Show this message") do
      puts opt
      exit
    end
    opt.on_tail('-v', '--version', "Show version") do
      puts defined?(SCRIPT_VERSION) ? SCRIPT_VERSION : '1.0'
      exit
    end

    unless options[:host]
      unless ARGV.empty?
        options[:host] = ARGV.first.split(':').first
      end
    end
    unless options[:port]
      if options[:host]
        options[:port] = options[:host] =~ /\:/ ? options[:host].split(':').last : 443
      end
    end
    begin
      opt.parse!(ARGV)
    rescue OptionParser::InvalidOption => e
      puts opt
      exit
    end
  end

  options
end

def root?
  unless Process.uid == 0; puts 'You are not root.'; exit; end
  true
end

def create_keystore(host)
  cmd = "keytool -importcert -noprompt -file #{host}.crt -keystore #{host}.jks -storepass changeit -alias #{host}.crt"
  puts "  $ #{cmd}" if options[:debug]
  `#{cmd}`
end

def import_keystore(host)
  puts "Importing keystore #{host}.jks"
  cmd = "keytool -importkeystore -srckeystore #{host}.jks -destkeystore #{options[:java_keystore]} -srcstorepass changeit -deststorepass changeit -noprompt"
  puts "  $ #{cmd}" if options[:debug]
  `#{cmd}`
end

def delete_keystore(host)
  cmd = "rm #{host}.jks"
  puts "  $ #{cmd}" if options[:debug]
  `#{cmd}`
end

def fetch_cert(host, port)
  puts "Fetching certificate for https://#{host}/"
  cmd = "echo | openssl s_client -connect #{host}:#{port} 2>&1 | sed -ne '/-BEGIN CERTIFICATE-/,/-END CERTIFICATE-/p'"
  puts "  $ #{cmd}" if options[:debug]
  `#{cmd}` 
end

def delete_alias(host)
  puts "Deleting old alias from existing keystore file: #{options[:local_keystore]}"
  cmd = "keytool -delete -trustcacerts -alias #{host} -keystore #{options[:local_keystore]} -deststorepass changeit -noprompt"
  puts "  $ #{cmd}" if options[:debug]
  `#{cmd}`
end

def add_alias(host)
  puts "Adding new alias to existing keystore file: #{options[:local_keystore]}"
  cmd = "keytool -import -trustcacerts -alias #{host} -file #{host}.crt -keystore #{options[:local_keystore]} -deststorepass changeit -noprompt"
  puts "  $ #{cmd}" if options[:debug]
  `#{cmd}`
end

def confirm_alias(host)
  puts "Confirming alias is in keystore file: #{options[:local_keystore]}"
  if alias_exists? host 
    "Confirmed"
  else 
    "Failed to add alias to #{options[:local_keystore]}"
  end
end

def alias_exists?(host)
  cmd = "keytool -list -keystore #{options[:local_keystore]} -deststorepass changeit | grep #{host}"
  puts "  $ #{cmd}" if options[:debug]
  a = `#{cmd}`.strip
  return !(a.empty?)
end

Main.new if __FILE__ == $0

